#include "..\CUDA_ComputeStructureTensor.h"
#include "CUDA_ST_Method_EuclideanVector.h"


inline __host__ __device__ float3 operator-(float3 a, float3 b)
{
	return make_float3(a.x - b.x, a.y - b.y, a.z - b.z);
}

inline __host__ __device__ float3 operator+(float3 a, float3 b)
{
	return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
}

inline __host__ __device__ float3 operator*(float3 a, const float b)
{
	return make_float3(a.x * b, a.y * b, a.z * b);
}

__device__ float CUDA_Point_Distance(const float3 *src, const float3 *dst)
{
	float x = src->x - dst->x;
	float y = src->y - dst->y;
	float z = src->z - dst->z;
	float d = sqrtf(x * x + y * y + z * z);
	return d;
}

__device__ float CUDA_DOT(const float3 *op1, const float3 *op2)
{
	float res = op1->x * op2->x + op1->y * op2->y + op1->z * op2->z;
	return res;
}

__device__ void CUDA_EigenSolve(const float *m, float *eigen_values, float *eigen_vectors)
{
	float t[6];
	t[0] = m[0]; t[1] = m[1]; t[2] = m[2]; t[3] = m[4]; t[4] = m[5]; t[5] = m[8];

	float mean = (t[0] + t[1] + t[2] + t[3] + t[4] + t[5]) / 6.0f;
	if (fabs(mean) < 10e-10) {
		eigen_values[0] = 0.0f;
		eigen_values[1] = 0.0f;
		eigen_values[2] = 0.0f;
		eigen_vectors[0] = eigen_vectors[1] = eigen_vectors[2] = 0.0f;
		eigen_vectors[3] = eigen_vectors[4] = eigen_vectors[5] = 0.0f;
		eigen_vectors[6] = eigen_vectors[7] = eigen_vectors[8] = 0.0f;
		return;
	}

	t[0] = t[0] / mean;
	t[1] = t[1] / mean;
	t[2] = t[2] / mean;
	t[3] = t[3] / mean;
	t[4] = t[4] / mean;
	t[5] = t[5] / mean;

	

	// comput the eigen values
	float I1 = t[0] + t[3] + t[5];
	float I2 = t[0] * t[3] + t[0] * t[5] + t[3] * t[5] - 
		(t[1] * t[1] + t[2] * t[2] + t[4] * t[4]);
	float I3 = t[0] * t[3] * t[5] + 2.0f * (t[1] * t[2] * t[4]) - 
		(t[0] * t[4] * t[4] + t[3] * t[2] * t[2] + t[5] * t[1] * t[1]);

	float PIthird = 3.1415926535f / 3.0f;
	float third = 1.0f / 3.0f;
	float I1third = I1 * third;
	float I1thirdsqr = I1third * I1third;
	float I2third = I2 * third;
	float v = I1thirdsqr - I2third;

	float s = I1thirdsqr * I1third - I1 * I2 / 6.0f + 0.5f * I3;

	if (v < 0.0f)
		v = 0.0f;
	float sqrtv = sqrtf(v);
	float temp = s / (v * sqrtv);
	if (temp > 1.0f)
		temp = 1.0f;
	if (temp < -1.0f)
		temp = -1.0f;
	float phi = acosf(temp) * third;
	float sqrtv2 = 2.0f * sqrtv;

	eigen_values[0] = I1third + sqrtv2 * cosf(phi);
	eigen_values[1] = I1third - sqrtv2 * cosf(PIthird+phi);
	eigen_values[2] = I1third - sqrtv2 * cosf(PIthird-phi);

	/* compute eigen vectors */
	eigen_vectors[0] = (t[1] * t[4] - (t[3] - eigen_values[0]) * t[2]) * (t[2] * t[4] - (t[5] - eigen_values[0]) * t[1]);
	eigen_vectors[1] = (t[4] * t[2] - (t[5] - eigen_values[0]) * t[1]) * (t[1] * t[2] - (t[0] - eigen_values[0]) * t[4]);
	eigen_vectors[2] = (t[4] * t[1] - (t[3] - eigen_values[0]) * t[2]) * (t[2] * t[1] - (t[0] - eigen_values[0]) * t[4]);

	eigen_vectors[3] = (t[1] * t[4] - (t[3] - eigen_values[1]) * t[2]) * (t[2] * t[4] - (t[5] - eigen_values[1]) * t[1]);
	eigen_vectors[4] = (t[4] * t[2] - (t[5] - eigen_values[1]) * t[1]) * (t[1] * t[2] - (t[0] - eigen_values[1]) * t[4]);
	eigen_vectors[5] = (t[4] * t[1] - (t[3] - eigen_values[1]) * t[2]) * (t[2] * t[1] - (t[0] - eigen_values[1]) * t[4]);

	eigen_vectors[6] = (t[1] * t[4] - (t[3] - eigen_values[2]) * t[2]) * (t[2] * t[4] - (t[5] - eigen_values[2]) * t[1]);
	eigen_vectors[7] = (t[4] * t[2] - (t[5] - eigen_values[2]) * t[1]) * (t[1] * t[2] - (t[0] - eigen_values[2]) * t[4]);
	eigen_vectors[8] = (t[4] * t[1] - (t[3] - eigen_values[2]) * t[2]) * (t[2] * t[1] - (t[0] - eigen_values[2]) * t[4]);

	float len = sqrtf(eigen_vectors[0] * eigen_vectors[0] + eigen_vectors[1] * eigen_vectors[1] + eigen_vectors[2] * eigen_vectors[2]);
	if (len > 0.0f) {
		eigen_vectors[0] = eigen_vectors[0] / len;
		eigen_vectors[1] = eigen_vectors[1] / len;
		eigen_vectors[2] = eigen_vectors[2] / len;
	}

	len = sqrtf(eigen_vectors[3] * eigen_vectors[3] + eigen_vectors[4] * eigen_vectors[4] + eigen_vectors[5] * eigen_vectors[5]);
	if (len > 0.0f) {
		eigen_vectors[3] = eigen_vectors[3] / len;
		eigen_vectors[4] = eigen_vectors[4] / len;
		eigen_vectors[5] = eigen_vectors[5] / len;
	}

	len = sqrtf(eigen_vectors[6] * eigen_vectors[6] + eigen_vectors[7] * eigen_vectors[7] + eigen_vectors[8] * eigen_vectors[8]);
	if (len > 0.0f) {
		eigen_vectors[6] = eigen_vectors[6] / len;
		eigen_vectors[7] = eigen_vectors[7] / len;
		eigen_vectors[8] = eigen_vectors[8] / len;
	}
}

__device__ void CUDA_SquareRoot(float *in, float *out)
{
	//out[0] = in[0];
	//out[1] = in[1];
	//out[2] = in[2];
	//out[3] = in[3];
	//out[4] = in[4];
	//out[5] = in[5];

	float m[9], va[3], ve[9];
	m[0] = in[0]; m[1] = in[1]; m[2] = in[2]; 
	m[3] = in[1]; m[4] = in[3]; m[5] = in[4];
	m[6] = in[2]; m[7] = in[4]; m[8] = in[5];

	CUDA_EigenSolve(m, va, ve);

	for (int i = 0; i < 3; ++i) {
		if (va[i] > 0.0f)
			va[i] = sqrt(va[i]);
		else
			va[i] = 0.0f;
	}

	float temp[9];
	temp[0] = va[0] * ve[0];
	temp[1] = va[1] * ve[3];
	temp[2] = va[2] * ve[6];
	temp[3] = va[0] * ve[1];
	temp[4] = va[1] * ve[4];
	temp[5] = va[2] * ve[7];
	temp[6] = va[0] * ve[2];
	temp[7] = va[1] * ve[5];
	temp[8] = va[2] * ve[8];

	out[0] = temp[0] * ve[0] + temp[1] * ve[3] + temp[2] * ve[6];
	out[1] = temp[0] * ve[1] + temp[1] * ve[4] + temp[2] * ve[7];
	out[2] = temp[0] * ve[2] + temp[1] * ve[5] + temp[2] * ve[8];
	out[3] = temp[3] * ve[1] + temp[4] * ve[4] + temp[5] * ve[7];
	out[4] = temp[3] * ve[2] + temp[4] * ve[5] + temp[5] * ve[8];
	out[5] = temp[6] * ve[2] + temp[7] * ve[5] + temp[8] * ve[8];
}

__device__ void CUDA_ComputeEuclideanVector(const CUDA_Fiber *fiber, const int length, float *vec)
{
	int F = fiber->fc > length ? length : fiber->fc;
	int B = fiber->bc > length ? length : fiber->bc;
	
	for (int i = 0; i < 9; ++i)
		vec[i] = 0.0f;

	if (F + B <= 2) {
		return;
	}

	float ss = 1.0f / float(F + B - 2);
	/* compute the mean vector */
	float3 mean = make_float3(0.0f, 0.0f, 0.0f);
	for (int k = F-1; k > 0; --k) {
		float3 temp = fiber->f[k-1] - fiber->f[k];
		temp = temp;
		mean = mean + temp;
	}
	for (int k = 0; k < B-1; ++k) {
		float3 temp = fiber->b[k+1] - fiber->b[k];
		temp = temp;
		mean = mean + temp;
	}
	mean = mean * ss;

	/* compute the Covariance Matrix */
	float cov[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
	for (int k = F-1; k > 0; --k) {
		float3 temp = fiber->f[k-1] - fiber->f[k];
		temp = temp;
		float3 d = temp - mean;
		cov[0] += d.x * d.x;	cov[1] += d.x * d.y;	cov[2] += d.x * d.z;
		cov[3] += d.y * d.y;	cov[4] += d.y * d.z;	cov[5] += d.z * d.z;
	}
	for (int k = 0; k < B-1; ++k) {
		float3 temp = fiber->b[k+1] - fiber->b[k];
		temp = temp;
		float3 d = temp - mean;
		cov[0] += d.x * d.x;	cov[1] += d.x * d.y;	cov[2] += d.x * d.z;
		cov[3] += d.y * d.y;	cov[4] += d.y * d.z;	cov[5] += d.z * d.z;
	}
	for (int i = 0; i < 6; ++i)
		cov[i] *= ss;
	
	vec[0] = mean.x;	vec[1] = mean.y;	vec[2] = mean.z;
	CUDA_SquareRoot(cov, &(vec[3]));
}

__global__ void 
FiberParameterization_EV_kernel(CUDA_Fiber *fibers, Fiber_EV_Param *parameters, const int length, const int num)
{
	int x = blockIdx.x * blockDim.x + threadIdx.x;
	int y = blockIdx.y * blockDim.y + threadIdx.y;

	int index = y * gridDim.x * blockDim.x + x;

	if (index < num) {
		/* copy the seed */
		parameters[index].seed.pos = fibers[index].seed.pos;
		parameters[index].seed.dir = fibers[index].seed.dir;

		/* parameterization */
		CUDA_ComputeEuclideanVector(&(fibers[index]), length, parameters[index].parameters);
	}
}


extern "C"
void CUDA_FiberParameterization_EV_DTI(CFiber **fibers,
		Fiber_EV_Param *parameters, const int w, const int h, const int length)
{
	int remain_num, last_num, this_num, num_per_run, i, j;
	
	remain_num = w * h;
	last_num = 0;
	num_per_run = THREADS_PER_RUN;

	CUDA_Fiber *h_fibers = (CUDA_Fiber *)malloc(sizeof(CUDA_Fiber)*num_per_run);
	
	CUDA_Fiber *d_fibers = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_fibers, sizeof(CUDA_Fiber)*num_per_run));
	Fiber_EV_Param *d_parameters = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_parameters, sizeof(Fiber_EV_Param)*num_per_run));

	while (remain_num > 0) {
		this_num = num_per_run < remain_num ? num_per_run : remain_num;
		for (i = 0; i < this_num; ++i) {
			h_fibers[i].fc = fibers[i+last_num][0].m_fCount;
			h_fibers[i].bc = fibers[i+last_num][0].m_bCount;
			h_fibers[i].seed.pos.x = fibers[i+last_num][0].m_seed.pos.m_x;
			h_fibers[i].seed.pos.y = fibers[i+last_num][0].m_seed.pos.m_y;
			h_fibers[i].seed.pos.z = fibers[i+last_num][0].m_seed.pos.m_z;
			h_fibers[i].seed.dir.x = fibers[i+last_num][0].m_seed.dir.m_x;
			h_fibers[i].seed.dir.y = fibers[i+last_num][0].m_seed.dir.m_y;
			h_fibers[i].seed.dir.z = fibers[i+last_num][0].m_seed.dir.m_z;

			for (j = 0; j < h_fibers[i].fc; ++j) {
				h_fibers[i].f[j].x = fibers[i+last_num][0].m_pF[j].m_x;
				h_fibers[i].f[j].y = fibers[i+last_num][0].m_pF[j].m_y;
				h_fibers[i].f[j].z = fibers[i+last_num][0].m_pF[j].m_z;
			}

			for (j = 0; j < h_fibers[i].bc; ++j) {
				h_fibers[i].b[j].x = fibers[i+last_num][0].m_pB[j].m_x;
				h_fibers[i].b[j].y = fibers[i+last_num][0].m_pB[j].m_y;
				h_fibers[i].b[j].z = fibers[i+last_num][0].m_pB[j].m_z;
			}
		}

		cutilSafeCall(cudaMemcpy(d_fibers, h_fibers, 
			sizeof(CUDA_Fiber)*this_num, cudaMemcpyHostToDevice));

		dim3 dimBlock(64, 1);
		dim3 dimGrid(num_per_run/64, 1);

		FiberParameterization_EV_kernel<<<dimGrid, dimBlock>>>(d_fibers, d_parameters, length, this_num);
		cutilSafeCall(cudaGetLastError());

		cutilSafeCall(cudaMemcpy(&(parameters[last_num]), d_parameters, 
			sizeof(Fiber_EV_Param)*this_num, cudaMemcpyDeviceToHost));

		remain_num -= this_num;
		last_num += this_num;

		printf("#");
		fflush(stdout);
	}

	cutilSafeCall(cudaFree(d_fibers));
	cutilSafeCall(cudaFree(d_parameters));

	free(h_fibers);
}

extern "C"
void CUDA_FiberParameterization_EV_HARDI(CFiber **fibers, int *count, 
		Fiber_EV_Param *parameters, const int w, const int h, const int length)
{
	int remain_num, last_num, this_num, num_per_run, i, j, k;
	
	remain_num = w * h * 3;
	last_num = 0;
	num_per_run = THREADS_PER_RUN * 3;

	CUDA_Fiber *h_fibers = (CUDA_Fiber *)malloc(sizeof(CUDA_Fiber)*num_per_run);
	memset(h_fibers, 0, sizeof(CUDA_Fiber)*num_per_run);
	
	CUDA_Fiber *d_fibers = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_fibers, sizeof(CUDA_Fiber)*num_per_run));
	Fiber_EV_Param *d_parameters = NULL;
	cutilSafeCall(cudaMalloc((void **)&d_parameters, sizeof(Fiber_EV_Param)*num_per_run));

	while (remain_num > 0) {
		this_num = num_per_run < remain_num ? num_per_run : remain_num;
		for (i = 0; i < this_num/3; ++i) {
			for (k = 0; k < count[i+last_num]; ++k) {
				h_fibers[i*3+k].fc = fibers[i+last_num][k].m_fCount;
				h_fibers[i*3+k].bc = fibers[i+last_num][k].m_bCount;
				h_fibers[i*3+k].seed.pos.x = fibers[i+last_num][k].m_seed.pos.m_x;
				h_fibers[i*3+k].seed.pos.y = fibers[i+last_num][k].m_seed.pos.m_y;
				h_fibers[i*3+k].seed.pos.z = fibers[i+last_num][k].m_seed.pos.m_z;
				h_fibers[i*3+k].seed.dir.x = fibers[i+last_num][k].m_seed.dir.m_x;
				h_fibers[i*3+k].seed.dir.y = fibers[i+last_num][k].m_seed.dir.m_y;
				h_fibers[i*3+k].seed.dir.z = fibers[i+last_num][k].m_seed.dir.m_z;

				for (j = 0; j < h_fibers[i*3+k].fc; ++j) {
					h_fibers[i*3+k].f[j].x = fibers[i+last_num][k].m_pF[j].m_x;
					h_fibers[i*3+k].f[j].y = fibers[i+last_num][k].m_pF[j].m_y;
					h_fibers[i*3+k].f[j].z = fibers[i+last_num][k].m_pF[j].m_z;
				}

				for (j = 0; j < h_fibers[i*3+k].bc; ++j) {
					h_fibers[i*3+k].b[j].x = fibers[i+last_num][k].m_pB[j].m_x;
					h_fibers[i*3+k].b[j].y = fibers[i+last_num][k].m_pB[j].m_y;
					h_fibers[i*3+k].b[j].z = fibers[i+last_num][k].m_pB[j].m_z;
				}
			}
		}

		cutilSafeCall(cudaMemcpy(d_fibers, h_fibers, 
			sizeof(CUDA_Fiber)*this_num, cudaMemcpyHostToDevice));

		dim3 dimBlock(64, 1);
		dim3 dimGrid(num_per_run/64, 1);

		FiberParameterization_EV_kernel<<<dimGrid, dimBlock>>>(d_fibers, d_parameters, length, this_num);
		cutilSafeCall(cudaGetLastError());

		cutilSafeCall(cudaMemcpy(&(parameters[last_num*3]), d_parameters, 
			sizeof(Fiber_EV_Param)*this_num, cudaMemcpyDeviceToHost));

		remain_num -= this_num;
		last_num += this_num/3;

		printf("#");
		fflush(stdout);
	}

	cutilSafeCall(cudaFree(d_fibers));
	cutilSafeCall(cudaFree(d_parameters));

	free(h_fibers);
}